
`%||%` <- function(a, b) if (is.null(a)) b else a

#' Plot values for a protein/gene of interest
#'
#' \code{plot_single} generates a barplot of a protein/gene of interest.
#'
#' @param object SummarizedExperiment,
#' Data object for which differentially enriched proteins are annotated
#' (output from \code{\link{test_diff}()} and \code{\link{add_rejections}()}).
#' @param proteins Character,
#' The name(s) of the protein(s)/gene(s) to plot.
#' @param type 'contrast' or 'centered',
#' The type of data scaling used for plotting.
#' Either the fold change ('contrast') or
#' the centered log2-intensity ('centered').
#' @param plot Logical(1),
#' If \code{TRUE} (default) the barplot is produced.
#' Otherwise (if \code{FALSE}), the summaries which the
#' barplot is based on are returned.
#' @return A barplot (generated by \code{\link[ggplot2]{ggplot}}).
#'
#' @examples
#' # Load example
#' data(Silicosis_pg)
#' data <- Silicosis_pg
#' data_unique <- make_unique(data, "Gene.names", "Protein.IDs", delim = ";")
#'
#' # Make SummarizedExperiment
#' ecols <- grep("LFQ.", colnames(data_unique))
#'
#'
#' ## Load experiement design
#' data(Silicosis_ExpDesign)
#' exp_design <- Silicosis_ExpDesign
#' se <- make_se(data_unique, ecols, exp_design)
#'
#' # Filter and normalize
#' filt <- filter_se(se, thr = 0, fraction = 0.4, filter_formula = ~ Reverse != "+" & Potential.contaminant!="+")
#'
#' norm <- normalize_vsn(filt)
#'
#' # Impute missing values using different functions
#' imputed <- impute(norm, fun = "MinProb", q = 0.05)
#'
#' # Test for differentially expressed proteins
#' diff <- test_diff(imputed, "control", "PBS")
#' dep <- add_rejections(diff, alpha = 0.05, lfc = 1)
#'
#' # Plot single proteins
#' plot_single(dep, 'Oas2')
#' plot_single(dep, 'Oas2', 'centered')
#' plot_single(dep, c('Oas2', 'Oas3'))
#' plot_single(dep, c('Oas2', 'Oas3'), plot = FALSE)
#' @export
plot_single <- function(object, proteins, type = c("contrast", "centered"), plot = TRUE) {
  assertthat::assert_that(class(object) %in% c("SummarizedExperiment",
                                               "DEGdata"), is.character(proteins), is.character(type),
                          is.logical(plot), length(plot) == 1)
  type <- match.arg(type)
  row_data <- rowData(object, use.names = FALSE)
  if (any(!c("label", "condition", "replicate") %in% colnames(colData(object)))) {
    stop("'label', 'condition' and/or 'replicate' columns are not present in '",
         objectarse(substitute(object)), "'\nRun make_se() or make_se_parse() to obtain the required columns",
         call. = FALSE)
  }
  if (length(grep("_p.adj|_diff", colnames(row_data))) < 1) {
    stop("'[contrast]_diff' and '[contrast]_p.adj' columns are not present in '",
         objectarse(substitute(object)), "'\nRun test_diff() to obtain the required columns",
         call. = FALSE)
  }
  if (!"name" %in% colnames(row_data)) {
    stop("'name' column not present in '", objectarse(substitute(object)),
         "'\nRun make_se() or make_se_parse() to obtain the required columns",
         call. = FALSE)
  }
  if (all(!proteins %in% row_data$name)) {
    if (length(proteins) == 1) {
      rows <- grep(substr(proteins, 1, nchar(proteins) -
                            1), row_data$name)
      possibilities <- row_data$name[rows]
    }
    else {
      rows <- lapply(proteins, function(x) grep(substr(x,
                                                       1, nchar(x) - 1), row_data$name))
      possibilities <- row_data$name[unlist(rows)]
    }
    if (length(possibilities) > 0) {
      possibilities_msg <- paste0("Do you mean: '", paste0(possibilities,
                                                           collapse = "', '"), "'")
    }
    else {
      possibilities_msg <- NULL
    }
    stop("please run `plot_single()` with a valid protein names in the 'proteins' argument\n",
         possibilities_msg, call. = FALSE)
  }
  if (any(!proteins %in% row_data$name)) {
    proteins <- proteins[proteins %in% row_data$name]
    warning("Only used the following protein(s): '", paste0(proteins,
                                                            collapse = "', '"), "'")
  }
  subset <- object[proteins]
  if (type == "centered") {
    means <- rowMeans(assay(subset), na.rm = TRUE)
    df_reps <- data.frame(assay(subset) - means) %>% rownames_to_column() %>%
      gather(ID, val, -rowname) %>% left_join(., data.frame(colData(subset)),
                                              by = "ID")
    df_reps$replicate <- as.factor(df_reps$replicate)
    df <- df_reps %>% group_by(condition, rowname) %>% summarize(mean = mean(val,
                                                                             na.rm = TRUE), sd = sd(val, na.rm = TRUE), n = n()) %>%
      mutate(error = qnorm(0.975) * sd/sqrt(n), CI.L = mean -
               error, CI.R = mean + error) %>% as.data.frame()
    df$rowname <- readr::parse_factor(df$rowname, levels = proteins)
    p <- ggplot(df, aes(condition, mean)) + geom_hline(yintercept = 0) +
      geom_col(colour = "black", fill = "grey") + geom_point(data = df_reps,
                                                             aes(condition, val, col = replicate), shape = 18,
                                                             size = 5, position = position_dodge(width = 0.3)) +
      geom_errorbar(aes(ymin = CI.L, ymax = CI.R), width = 0.3) +
      labs(x = "Baits", y = expression(log[2] ~ "Centered intensity" ~
                                         "(±95% CI)"), col = "Rep") + facet_wrap(~rowname) +
      theme_DEP1()
  }
  if (type == "contrast") {
    df <- rowData(subset, use.names = FALSE) %>% data.frame() %>%
      select(name, ends_with("_diff"), ends_with("_CI.L"),
             ends_with("_CI.R")) %>% gather(var, val, -name) %>%
      mutate(contrast = gsub("_diff|_CI.L|_CI.R", "", var),
             var = gsub(".*_", "", var)) %>% spread(var, val)
    df$name <- readr::parse_factor(df$name, levels = proteins)
    suffix <- get_suffix(df$contrast)
    if (length(suffix)) {
      if (length(grep("_vs_", suffix)) > 0)
        suffix = str_extract(suffix, "_vs_(.)+")
      df$contrast <- gsub(suffix, "", df$contrast)
    }
    p <- ggplot(df, aes(contrast, diff)) + geom_hline(yintercept = 0) +
      geom_col(colour = "black", fill = "grey") + geom_errorbar(aes(ymin = CI.L,
                                                                    ymax = CI.R), width = 0.3) + labs(x = suffix, y = expression(log[2] ~
                                                                                                                                   "Fold change" ~ "(±95% CI)")) + facet_wrap(~name) +
      DEP2::theme_DEP1()
  }
  if (plot) {
    return(p)
  }
  else {
    if (type == "centered") {
      df <- df %>% select(rowname, condition, mean, CI.L,
                          CI.R)
      colnames(df) <- c("protein", "condition", "log2_intensity",
                        "CI.L", "CI.R")
    }
    if (type == "contrast") {
      df <- df %>% select(name, contrast, diff, CI.L, CI.R) %>%
        mutate(contrast = paste0(contrast, suffix))
      colnames(df) <- c("protein", "contrast", "log2_fold_change",
                        "CI.L", "CI.R")
    }
    return(df)
  }
}


# dds; genes = rownames(dds_res)[1]; intgroup = "condition"; annotation_obj = NULL;
# transform = TRUE; ylimZero = TRUE; add_labels = F; labels_repel = T
# plot_single_deg(dds_res,genes = rownames(dds_res)[1], type = "contrast")

plot_single_deg <- function(deg, genes, type = c("contrast", "centered"), plot = TRUE) {
  # Show error if inputs are not the required classes
  assertthat::assert_that(inherits(deg, "DEGdata"),
                          is.character(genes),
                          is.character(type),
                          is.logical(plot),
                          length(plot) == 1)

  # Show error if inputs do not contain required columns
  type <- match.arg(type)

  row_data <- rowData(deg, use.names = T)


  if(length(grep("_p.adj|_diff", colnames(row_data))) < 1) {
    stop("'[contrast]_diff' and '[contrast]_p.adj' columns are not present in '",
         deparse(substitute(deg)),
         "'\nRun test_diff_deg() to obtain the required columns",
         call. = FALSE)
  }
  row_data$name <- rownames(row_data)
  # Show error if an unvalid protein name is given
  if(all(!genes %in% row_data$name)) {
    if(length(genes) == 1) {
      rows <- grep(substr(genes, 1, nchar(genes) - 1),row_data$name)
      possibilities <- row_data$name[rows]
    } else {
      rows <- lapply(genes, function(x)
        grep(substr(x, 1, nchar(x) - 1),row_data$name))
      possibilities <- row_data$name[unlist(rows)]
    }

    if(length(possibilities) > 0) {
      possibilities_msg <- paste0(
        "Do you mean: '",
        paste0(possibilities, collapse = "', '"),
        "'")
    } else {
      possibilities_msg <- NULL
    }
    stop("please run `plot_single_deg()` with a valid genes in the 'genes' argument\n",
         possibilities_msg,
         call. = FALSE)
  }
  if(any(!genes %in% row_data$name)) {
    genes <- genes[genes %in% row_data$name]
    warning("Only used the following protein(s): '",
            paste0(genes, collapse = "', '"),
            "'")
  }

  # Single protein
  subset <- deg[genes]

  # Plot either the centered log-intensity values
  # per condition ('centered') or the average fold change of conditions
  # versus the control condition ('contrast') for a single protein
  if(type == "centered") {
    # Obtain protein-centered fold change values
    counts_assay = subset@ntf
    means <- rowMeans(counts_assay, na.rm = TRUE)
    df_reps <- data.frame(counts_assay - means) %>%
      rownames_to_column() %>%
      gather(ID, val, -rowname) %>%
      left_join(., data.frame(colData(subset)), by = "ID")
    df_reps$replicate <- as.factor(df_reps$replicate)
    df <- df_reps %>%
      group_by(condition, rowname) %>%
      summarize(mean = mean(val, na.rm = TRUE),
                sd = sd(val, na.rm = TRUE),
                n = n()) %>%
      mutate(error = qnorm(0.975) * sd / sqrt(n),
             CI.L = mean - error,
             CI.R = mean + error) %>%
      as.data.frame()
    df$rowname <- readr::parse_factor(df$rowname, levels = genes)

    # Plot the centered intensity values for the replicates and the mean
    p <- ggplot(df, aes(condition, mean)) +
      geom_hline(yintercept = 0) +
      geom_col(colour = "black", fill = "grey") +
      geom_point(data = df_reps, aes(condition, val, col = replicate),
                 shape = 18, size = 5, position = position_dodge(width=0.3)) +
      geom_errorbar(aes(ymin = CI.L, ymax = CI.R), width = 0.3) +
      labs(x = "Baits",
           y = expression(log[2]~"Centered intensity"~"(\u00B195% CI)"),
           col = "Rep") +
      facet_wrap(~rowname) +
      theme_DEP1()
  }
  if(type == "contrast") {
    # Select values for a single protein
    df <- rowData(subset, use.names = T) %>%
      data.frame() %>%
      rownames_to_column(., "name") %>%
      # data.frame() %>%
      select(name,
             ends_with("_diff")
             ) %>%
      gather(var, val, -name) %>%
      mutate(contrast = gsub("_diff|_CI.L|_CI.R", "", var),
             var = gsub(".*_", "", var)) %>%
      spread(var, val)
    df$name <- readr::parse_factor(df$name, levels = genes)
    if(nrow(df) > 1){
      suffix <- get_suffix(df$contrast)
    }else{
      suffix <- NULL
    }

    if(length(suffix)) {
      if(length(grep("_vs_",suffix)) > 0)
        suffix = str_extract(suffix, "_vs_(.)+")
      df$contrast <- gsub(suffix, "", df$contrast)
    }
    # Plot the average fold change of conditions versus the control condition
    p <- ggplot(df, aes(contrast, diff)) +
      geom_hline(yintercept = 0) +
      geom_col(colour = "black", fill = "grey") +
      # geom_errorbar(aes(ymin = CI.L, ymax = CI.R), width = 0.3) +
      labs(x = suffix,
           y = expression(log[2]~"Fold change"~"(\u00B195% CI)")) +
      facet_wrap(~name) +
      theme_DEP1()
  }
  if(plot) {
    return(p)
  } else {
    if(type == "centered") {
      df <- df %>%
        select(rowname, condition, mean, CI.L, CI.R)
      colnames(df) <- c("protein", "condition",
                        "log2_intensity", "CI.L", "CI.R")
    }
    if(type == "contrast") {
      df <- df %>%
        select(name, contrast, diff, CI.L, CI.R) %>%
        mutate(contrast = paste0(contrast, suffix))
      colnames(df) <- c("protein", "contrast",
                        "log2_fold_change", "CI.L", "CI.R")
    }
    return(df)
  }

}

## Plot heatmap
setGeneric("plot_heatmap", function(object,
                                    type = c("centered","contrast"),
                                    manual_contrast = NULL,
                                    kmeans = FALSE, k = 6,
                                    color = c("RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral"), col_limit = 6,
                                    indicate = NULL,
                                    row_font_size = 6, col_font_size = 10,
                                    clustering_distance = c("euclidean", "maximum", "manhattan", "canberra", "binary",
                                                            "minkowski", "pearson", "spearman", "kendall", "gower"),
                                    split_order = NULL,
                                    label_few_peptide_rows = FALSE,
                                    chooseToshow = NULL,
                                    plot = TRUE, seed = 42,
                                    ...) {
  standardGeneric("plot_heatmap")
})


#' Plot a heatmap
#'
#' \code{plot_heatmap} generates a heatmap of all or selected features for DEP or DEG data.
#'
#' @param object SummarizedExperiment or DEGDdata
#' object for which differentially enriched proteins are annotated
#' (output from \code{\link{test_diff}()} and \code{\link{add_rejections}()}).
#' @param type type 'contrast' or 'centered',
#' The type of data scaling used for plotting.
#' Either the fold change ('contrast') or
#' the centered log2-intensity ('centered').
#' @param manual_contrast Character.
#' Only plot significant proteins in certain contrasts.
#' @param kmeans Logical(1),
#' Whether or not to perform k-means clustering.
#' @param k Integer(1),
#' Sets the number of k-means clusters, k must less than row of heatmap
#' @param clustering_distance Character(1)."euclidean", "maximum", "manhattan", "canberra",
#' "binary", "minkowski", "pearson", "spearman", "kendall" or "gower",
#' Or a Function used to calculate clustering distance (for proteins and samples).
#' Based on \code{\link[ComplexHeatmap]{Heatmap}}
#' and \code{\link[cluster]{daisy}}.
#' @param color Character(1), sets the color panel (from \pkg{RColorBrewer}).
#' It can be a palette name in "RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral".
#' @param col_limit  Integer(1),
#' Sets the outer limits of the color scale.
#' @param indicate Character,
#' Sets additional annotation on the top of the heatmap
#' based on columns from the experimental design (colData).
#' Only applicable to type = 'centered'.
#' @param row_font_size Integer(1),
#' Font size of row name
#' @param col_font_size Integer(1),
#' Font size of columns label.
#' @param split_order a integer vector in 1:k.
#' The order of cluster slices in the heatmap. If NULL, slices order is based clustering distance. Work if Kmeans = TRUE.
#' @param label_few_peptide_rows Logic(1)
#' If TRUE, label proteins with few peptide by color.
#' Proteins with only one peptide show a purple name in heatmap, with two peptide show a blue name.
#' @param chooseToshow Character vector of the chosen proteins/peptides names.
#' If is not null, only plot the chosen proteins/peptides in it
#' @param plot Logical(1)
#' If \code{TRUE} (default) the heatmap is produced.
#' Otherwise (if \code{FALSE}), the data which the
#' heatmap is based on are returned.
#' @param seed Integer(1), the random seed that effect kmeans cluster in heatmap.
#' Preset a seed could a repeatable cluster result.
#' @param ... Other parameters to \code{\link[ComplexHeatmap]{Heatmap}}.
#'
#' @return
#' return a heatmap (generated by \code{\link[ComplexHeatmap]{Heatmap}}) or a data.frame if plot = F.
#' @examples
#'
#' # Load example
#' data(Silicosis_pg)
#' data <- Silicosis_pg
#' data_unique <- make_unique(data, "Gene.names", "Protein.IDs", delim = ";")
#'
#' # Differential test
#' ecols <- grep("LFQ.", colnames(data_unique))
#' se <- make_se_parse(data_unique, ecols,mode = "delim")
#' filt <- filter_se(se, thr = 0, fraction = 0.4, filter_formula = ~ Reverse != "+" & Potential.contaminant!="+")
#' norm <- normalize_vsn(filt)
#' imputed <- impute(norm, fun = "MinProb", q = 0.05)
#' diff <- test_diff(imputed, type = "control", control  = c("PBS"), fdr.type = "Storey's qvalue")
#' dep <- add_rejections(diff, alpha = 0.01,lfc = 2)
#'
#' # Heatmap
#' plot_heatmap(dep)
#'
#' @name plot_heatmap
#'
#' @rdname plot_heatmap
#' @exportMethod plot_heatmap
setMethod("plot_heatmap",
          "SummarizedExperiment",
          function(object, ...){
            .plot_heatmap.se(object = object,
                             type = type, manual_contrast = manual_contrast, kmeans = kmeans, k = k,
                             color = color, col_limit = col_limit, indicate = indicate, row_font_size = row_font_size,
                             col_font_size = col_font_size, clustering_distance = clustering_distance, split_order = split_order,
                             label_few_peptide_rows = label_few_peptide_rows, chooseToshow = chooseToshow, plot = plot,
                             seed = seed,
                             ...)
          }
)
.plot_heatmap.se <- function (object,
                             type = c("centered","contrast"),
                             manual_contrast = NULL,
                             # same_trend = FALSE,
                             kmeans = FALSE, k = 6,
                             color = c("RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral"), col_limit = 6,
                             indicate = NULL,
                             row_font_size = 6, col_font_size = 10,
                             clustering_distance = c("euclidean", "maximum", "manhattan", "canberra", "binary",
                                                     "minkowski", "pearson", "spearman", "kendall", "gower"),
                             split_order = NULL,
                             label_few_peptide_rows = FALSE,
                             # if_chooseToshow = FALSE,
                             chooseToshow = NULL,
                             plot = TRUE, seed = 42,...)
{
  # library(ComplexHeatmap)
  # library(tidyr)
  # library(tibble)
  if (is.integer(k))
    k <- as.numeric(k)
  if (is.integer(col_limit))
    col_limit <- as.numeric(col_limit)
  if (is.integer(row_font_size))
    row_font_size <- as.numeric(row_font_size)
  if (is.integer(col_font_size))
    col_font_size <- as.numeric(col_font_size)
  assertthat::assert_that(inherits(object, "SummarizedExperiment"),
                          is.character(type), is.logical(kmeans), is.numeric(k),
                          length(k) == 1, is.numeric(col_limit), length(col_limit) ==
                            1, is.numeric(row_font_size), length(row_font_size) ==
                            1, is.numeric(col_font_size), length(col_font_size) ==
                            1, is.logical(plot), length(plot) == 1)

  # set rowname of heatmap color, which Peptides == 1 to purple, ==2 to blue
  if(label_few_peptide_rows) {
    rowData(object)$row_name_color = "black"
    rowData(object)$row_name_color[which(rowData(object)$Peptides == 1)] = "purple"
    rowData(object)$row_name_color[which(rowData(object)$Peptides == 2)] = "blue"
  }

  type <- match.arg(type)
  clustering_distance <- match.arg(clustering_distance)
  color <- match.arg(color)

  if (!is.null(indicate) & type == "contrast") {
    warning("Heatmap annotation only applicable for type = 'centered'",
            call. = FALSE)
  }
  if (!is.null(indicate) & type == "centered") {
    ha1 <- get_annotation(object, indicate)
  } else {
    ha1 <- NULL
  }

  cat()
  ## exctract signicant protein
  filtered <- get_signicant(object, contrast = manual_contrast)
  colname <- colnames(assay(filtered))


  ## choose portion to plot
  if(!is.null(chooseToshow)){
    chooseToshow = intersect(chooseToshow,rownames(filtered))
    if(length(chooseToshow) == 0){
      stop("The chosen candidates are not significant")
    }else if(length(chooseToshow) <= k && kmeans){
      stop("The number of chosen candidates should larger than clusters number k.")
    }
    filtered <- filtered[match(chooseToshow,rownames(filtered)),]
  }

  # Check for missing values
  if(!is.null(manual_contrast)){
    ind = which(gsub("_[[:digit:]]*", "", colname) %in% unlist(strsplit(manual_contrast, split = "_vs_")))
    if(any(is.na(assay(filtered)[, ind]))) {
      warning("Missing values in '", deparse(substitute(object)), "'. ",
              "Using clustering_distance = 'gower'",
              call. = FALSE)
      clustering_distance <- "gower"
      obs_NA <- TRUE
    } else {
      obs_NA <- FALSE
    }

  } else{
    if(any(is.na(assay(filtered)))) {
      warning("Missing values in '", deparse(substitute(object)), "'. ",
              "Using clustering_distance = 'gower'",
              call. = FALSE)
      clustering_distance <- "gower"
      obs_NA <- TRUE
    } else {
      obs_NA <- FALSE
    }

  }

  # Get centered intensity values
  if(type == "centered") {
    if(!is.null(manual_contrast)){
      rowData(filtered)$mean <- rowMeans(assay(filtered)[ , ind] %>% as.matrix(), na.rm = TRUE)
      df <- assay(filtered)[ , ind] - rowData(filtered, use.names = FALSE)$mean
    } else{
      rowData(filtered)$mean <- rowMeans(assay(filtered), na.rm = TRUE)
      df <- assay(filtered) - rowData(filtered, use.names = FALSE)$mean
    }
  }

  # Get contrast fold changes ('contrast')
  if(type == "contrast") {
    df <- rowData(filtered, use.names = FALSE) %>%
      data.frame() %>%
      column_to_rownames(var = "name") %>%
      dplyr::select(ends_with("_diff"))
    colnames(df) <-
      gsub("_diff", "", colnames(df)) %>%
      gsub("_vs_", " vs ", .)
    df <- as.matrix(df)
    if(length(manual_contrast)>1){
      i = gsub("_vs_", " vs ", manual_contrast)
      ii = as.data.frame(df[ , i])
      colnames(ii) = i
      df = as.matrix(ii)
    }
  }

  # Facultative kmeans clustering
  if(kmeans & obs_NA) {
    warning("Cannot perform kmeans clustering with missing values",
            call. = FALSE)
    kmeans <- FALSE
  }
  if(kmeans & !obs_NA) {
    set.seed(seed)
    df_kmeans <- kmeans(df, k)
    if(type == "centered") {
      # Order the k-means clusters according to the maximum fold change
      # in all samples averaged over the proteins in the cluster
      order <- data.frame(df) %>%
        cbind(., cluster = df_kmeans$cluster) %>%
        mutate(row = apply(.[, seq_len(ncol(.) - 1)], 1, function(x) max(x))) %>%
        group_by(cluster) %>%
        summarize(index = sum(row)/n()) %>%
        arrange(desc(index)) %>%
        pull(cluster) %>%
        match(seq_len(k), .)
      df_kmeans$cluster <- order[df_kmeans$cluster]
    }
    if(type == "contrast") {
      # Order the k-means clusters according to their average fold change
      order <- data.frame(df) %>%
        cbind(df, cluster = df_kmeans$cluster) %>%
        gather(condition, diff, -cluster) %>%
        group_by(cluster) %>%
        summarize(row = mean(diff)) %>%
        arrange(desc(row)) %>%
        pull(cluster) %>%
        match(seq_len(k), .)
      df_kmeans$cluster <- order[df_kmeans$cluster]
    }
  }

  if(clustering_distance == "gower") {
    clustering_distance <- function(x) {
      dist <- cluster::daisy(x, metric = "gower")
      dist[is.na(dist)] <- max(dist, na.rm = TRUE)
      return(dist)
    }
  }

  # Legend info
  legend <- ifelse(type == "contrast",
                   "log2 Fold change",
                   "log2 Centered intensity")

  #set rowname color
  row_name_color = if(label_few_peptide_rows) {rowData(filtered)$row_name_color} else {"black"}

  # reorder the split
  if(!is.null(split_order)) {
    cluster_row_slices = FALSE
    cluster_column_slices = FALSE
    split_order = unique(split_order)
    split = if(kmeans) {
      df = df[which(df_kmeans$cluster %in% split_order),]
      factor(df_kmeans$cluster[df_kmeans$cluster %in% split_order], levels = split_order[split_order %in% df_kmeans$cluster])
    } else {NULL}
  } else {
    cluster_row_slices = TRUE
    cluster_column_slices = TRUE
    split = if(kmeans) {df_kmeans$cluster} else {NULL}
  }

  ht_df <- df

  # Heatmap
  # cat(color)
  # if(is.null(chooseToshow)) {
    ht1 = ComplexHeatmap::Heatmap(df,
                                  col = circlize::colorRamp2(
                                    seq(-col_limit, col_limit, (col_limit/5)),
                                    rev(RColorBrewer::brewer.pal(11, color))),
                                  split = split, # cluster by kmeans
                                  cluster_row_slices = cluster_row_slices,
                                  cluster_column_slices = cluster_column_slices,
                                  # cluster_rows = col_clust,
                                  # cluster_columns = row_clust,
                                  row_names_side = "left",
                                  column_names_side = "top",
                                  clustering_distance_rows = clustering_distance,
                                  clustering_distance_columns = clustering_distance,
                                  heatmap_legend_param = list(color_bar = "continuous",
                                                              legend_direction = "horizontal",
                                                              legend_width = unit(5, "cm"),
                                                              title_position = "lefttop"),
                                  name = legend,
                                  row_names_gp = gpar(fontsize = row_font_size, col = row_name_color),
                                  column_names_gp = gpar(fontsize = col_font_size),
                                  top_annotation = ha1,
                                  ...)

  # } else {
    # ht1 = ComplexHeatmap::Heatmap(df,
    #                               col = circlize::colorRamp2(
    #                                 seq(-col_limit, col_limit, (col_limit/5)),
    #                                 rev(RColorBrewer::brewer.pal(11, color))),
    #                               # split = split,
    #                               # cluster_row_slices = cluster_row_slices,
    #                               # cluster_column_slices = cluster_column_slices,
    #                               # cluster_rows = col_clust,
    #                               # cluster_columns = row_clust,
    #                               row_names_side = "left",
    #                               column_names_side = "top",
    #                               clustering_distance_rows = clustering_distance,
    #                               clustering_distance_columns = clustering_distance,
    #                               heatmap_legend_param = list(color_bar = "continuous",
    #                                                           legend_direction = "horizontal",
    #                                                           legend_width = unit(5, "cm"),
    #                                                           title_position = "lefttop"),
    #                               name = legend,
    #                               row_names_gp = gpar(fontsize = row_font_size, col = row_name_color),
    #                               column_names_gp = gpar(fontsize = col_font_size),
    #                               top_annotation = ha1,
    #                               ...)

  # }

  if(plot) {
    # Plot
    draw(ht1, heatmap_legend_side = "top")
  } else {
    # Return data.frame
    colnames(df) <- gsub(" ", "_", colnames(df))
    df <- df[, unlist(column_order(ht1))]
    if(kmeans) {
      df <- cbind(df, k = df_kmeans$cluster)
    }
    return <- df[unlist(row_order(ht1)),]
    data.frame(name = row.names(return), return) %>%
      mutate(order = row_number())
  }
}


#' @exportMethod plot_heatmap
#' @rdname plot_heatmap
setMethod("plot_heatmap",
          "DEGdata",
          function(object, ...){
            .plot_heatmap.DEGdata(object = object,
                                  type = type, manual_contrast = manual_contrast, kmeans = kmeans, k = k,
                                  color = color, col_limit = col_limit, indicate = indicate, row_font_size = row_font_size,
                                  col_font_size = col_font_size, clustering_distance = clustering_distance, split_order = split_order,
                                  chooseToshow = chooseToshow, plot = plot, seed = seed,
                                  ...)
          }
)
.plot_heatmap.DEGdata <- function (object,
                                  type = c("centered","contrast"),
                                  manual_contrast = NULL,
                                  # same_trend = FALSE,
                                  kmeans = FALSE, k = 6,
                                  color = c("RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral"),
                                  col_limit = 6,
                                  indicate = NULL,
                                  row_font_size = 6, col_font_size = 10,
                                  clustering_distance = c("euclidean", "maximum", "manhattan", "canberra", "binary",
                                                          "minkowski", "pearson", "spearman", "kendall", "gower"),
                                  split_order = NULL,
                                  # label_few_peptide_rows = FALSE,
                                  # if_chooseToshow = FALSE,
                                  chooseToshow = NULL,
                                  plot = TRUE, seed = 42, ...
                                  )
{

  if (is.integer(k))
    k <- as.numeric(k)
  if (is.integer(col_limit))
    col_limit <- as.numeric(col_limit)
  if (is.integer(row_font_size))
    row_font_size <- as.numeric(row_font_size)
  if (is.integer(col_font_size))
    col_font_size <- as.numeric(col_font_size)
  assertthat::assert_that(class(object) == "DEGdata",
                          is.character(type), is.logical(kmeans), is.numeric(k),
                          length(k) == 1, is.numeric(col_limit), length(col_limit) ==
                            1, is.numeric(row_font_size), length(row_font_size) ==
                            1, is.numeric(col_font_size), length(col_font_size) ==
                            1, is.logical(plot), length(plot) == 1)

  # # set rowname of heatmap color, which Peptides == 1 to purple, ==2 to blue
  # if(label_few_peptide_rows) {
  #   rowData(object)$row_name_color = "black"
  #   rowData(object)$row_name_color[which(rowData(object)$Peptides == 1)] = "purple"
  #   rowData(object)$row_name_color[which(rowData(object)$Peptides == 2)] = "blue"
  # }

  type <- match.arg(type)
  clustering_distance <- match.arg(clustering_distance)
  color <- match.arg(color)

  if (!is.null(indicate) & type == "contrast") {
    warning("Heatmap annotation only applicable for type = 'centered'",
            call. = FALSE)
  }
  if (!is.null(indicate) & type == "centered") {
    ha1 <- get_annotation(object, indicate)
  } else {
    ha1 <- NULL
  }

  message("deg heatmap")
  ## exctract signicant protein
  filtered <- get_signicant(object, contrast = manual_contrast, return_type = "subset")

  ## choose portion to plot
  if(!is.null(chooseToshow)){
    chooseToshow = intersect(chooseToshow,rownames(filtered))
    if(length(chooseToshow) == 0){
      warning("The chosen candidates are not significant")
    }
    filtered <- filtered[match(chooseToshow,rownames(filtered)),]
  }

  ntf = filtered@ntf
  rownames(ntf) = rownames(filtered)
  if(all(dim(ntf) == c(0,0))){
    stop("The ntf slot is empty, please run test_diff_deg, or assign it by a normalized assay")
  }
  if( any(dim(ntf) != dim(assay(filtered))) ){
    stop("The size of ntf is uncorrect.")
  }
  colname <- colnames(ntf)

  # Check for missing values
  if(!is.null(manual_contrast)){
    ind = which(gsub("_[[:digit:]]*", "", colname) %in% unlist(strsplit(manual_contrast, split = "_vs_")))
    if(any(is.na(ntf[, ind]))) {
      warning("Missing values in '", deparse(substitute(object)), "'. ",
              "Using clustering_distance = 'gower'",
              call. = FALSE)
      clustering_distance <- "gower"
      obs_NA <- TRUE
    } else {
      obs_NA <- FALSE
    }

  } else{
    if(any(is.na(ntf))) {
      warning("Missing values in '", deparse(substitute(object)), "'. ",
              "Using clustering_distance = 'gower'",
              call. = FALSE)
      clustering_distance <- "gower"
      obs_NA <- TRUE
    } else {
      obs_NA <- FALSE
    }

  }

  # Get centered intensity values
  if(type == "centered") {
    if(!is.null(manual_contrast)){
      rowData(filtered)$mean <- rowMeans(ntf[ , ind] %>% as.matrix(), na.rm = TRUE)
      df <- ntf[ , ind] - rowData(filtered, use.names = FALSE)$mean
    } else{
      rowData(filtered)$mean <- rowMeans(ntf, na.rm = TRUE)
      df <- ntf - rowData(filtered, use.names = FALSE)$mean
    }
  }

  # Get contrast fold changes ('contrast')
  if(type == "contrast") {
    df <- rowData(filtered, use.names = FALSE) %>%
      data.frame() %>%
      column_to_rownames(var = "name") %>%
      dplyr::select(ends_with("_diff"))
    colnames(df) <-
      gsub("_diff", "", colnames(df)) %>%
      gsub("_vs_", " vs ", .)
    df <- as.matrix(df)
    if(manual){
      i = gsub("_vs_", " vs ", manual_contrast)
      ii = as.data.frame(df[ , i])
      colnames(ii) = i
      df = as.matrix(ii)
    }
  }

  # Facultative kmeans clustering
  if(kmeans & obs_NA) {
    warning("Cannot perform kmeans clustering with missing values",
            call. = FALSE)
    kmeans <- FALSE
  }
  if(kmeans & !obs_NA) {
    set.seed(seed)
    df_kmeans <- kmeans(df, k)
    if(type == "centered") {
      # Order the k-means clusters according to the maximum fold change
      # in all samples averaged over the proteins in the cluster
      order <- data.frame(df) %>%
        cbind(., cluster = df_kmeans$cluster) %>%
        mutate(row = apply(.[, seq_len(ncol(.) - 1)], 1, function(x) max(x))) %>%
        group_by(cluster) %>%
        summarize(index = sum(row)/n()) %>%
        arrange(desc(index)) %>%
        pull(cluster) %>%
        match(seq_len(k), .)
      df_kmeans$cluster <- order[df_kmeans$cluster]
    }
    if(type == "contrast") {
      # Order the k-means clusters according to their average fold change
      order <- data.frame(df) %>%
        cbind(df, cluster = df_kmeans$cluster) %>%
        gather(condition, diff, -cluster) %>%
        group_by(cluster) %>%
        summarize(row = mean(diff)) %>%
        arrange(desc(row)) %>%
        pull(cluster) %>%
        match(seq_len(k), .)
      df_kmeans$cluster <- order[df_kmeans$cluster]
    }
  }

  if(clustering_distance == "gower") {
    clustering_distance <- function(x) {
      dist <- cluster::daisy(x, metric = "gower")
      dist[is.na(dist)] <- max(dist, na.rm = TRUE)
      return(dist)
    }
  }

  # Legend info
  legend <- ifelse(type == "contrast",
                   "log2 Fold change",
                   "log2 Centered intensity")

  #set rowname color
  row_name_color = "black"

  # reorder the split
  if(!is.null(split_order)) {
    cluster_row_slices = FALSE
    cluster_column_slices = FALSE
    split_order = unique(split_order)
    split = if(kmeans) {
      df = df[which(df_kmeans$cluster %in% split_order),]
      factor(df_kmeans$cluster[df_kmeans$cluster %in% split_order], levels = split_order[split_order %in% df_kmeans$cluster])
    } else {NULL}
  } else {
    cluster_row_slices = TRUE
    cluster_column_slices = TRUE
    split = if(kmeans) {df_kmeans$cluster} else {NULL}
  }

  ht_df <- df

  # Heatmap
  # cat(color)
  ht1 = ComplexHeatmap::Heatmap(df,
                                col = circlize::colorRamp2(
                                  seq(-col_limit, col_limit, (col_limit/5)),
                                  rev(RColorBrewer::brewer.pal(11, color))),
                                split = split, # cluster by kmeans
                                cluster_row_slices = cluster_row_slices,
                                cluster_column_slices = cluster_column_slices,
                                # cluster_rows = col_clust,
                                # cluster_columns = row_clust,
                                row_names_side = "left",
                                column_names_side = "top",
                                clustering_distance_rows = clustering_distance,
                                clustering_distance_columns = clustering_distance,
                                heatmap_legend_param = list(color_bar = "continuous",
                                                            legend_direction = "horizontal",
                                                            legend_width = unit(5, "cm"),
                                                            title_position = "lefttop"),
                                name = legend,
                                row_names_gp = gpar(fontsize = row_font_size, col = row_name_color),
                                column_names_gp = gpar(fontsize = col_font_size),
                                top_annotation = ha1,
                                ...)

  if(plot) {
    # Plot
    draw(ht1, heatmap_legend_side = "top")
  } else {
    # Return data.frame
    colnames(df) <- gsub(" ", "_", colnames(df))
    df <- df[, unlist(column_order(ht1))]
    if(kmeans) {
      df <- cbind(df, k = df_kmeans$cluster)
    }
    return <- df[unlist(row_order(ht1)),]
    data.frame(name = row.names(return), return) %>%
      mutate(order = row_number())
  }
}



# Internal function to get ComplexHeatmap::HeatmapAnnotation object
get_annotation <- function(dep, indicate) {
  assertthat::assert_that(
    inherits(dep, "SummarizedExperiment"),
    is.character(indicate))

  # Check indicate columns
  col_data <- colData(dep) %>%
    as.data.frame()
  columns <- colnames(col_data)
  if(all(!indicate %in% columns)) {
    stop("'",
         paste0(indicate, collapse = "' and/or '"),
         "' column(s) is/are not present in ",
         deparse(substitute(dep)),
         ".\nValid columns are: '",
         paste(columns, collapse = "', '"),
         "'.",
         call. = FALSE)
  }
  if(any(!indicate %in% columns)) {
    indicate <- indicate[indicate %in% columns]
    warning("Only used the following indicate column(s): '",
            paste0(indicate, collapse = "', '"),
            "'")
  }

  # Get annotation
  anno <- select(col_data, indicate)

  # Annotation color
  names <- colnames(anno)
  anno_col <- vector(mode="list", length=length(names))
  names(anno_col) <- names
  for(i in names) {
    var = anno[[i]] %>% unique() %>% sort()
    if(length(var) == 1)
      cols <- c("black")
    if(length(var) == 2)
      cols <- c("orangered", "cornflowerblue")
    if(length(var) <= 7 & length(var) > 2)
      cols <- RColorBrewer::brewer.pal(length(var), "Pastel1")
    if(length(var) > 7)
      cols <- RColorBrewer::brewer.pal(length(var), "Set3")
    names(cols) <- var
    anno_col[[i]] <-  cols
  }

  # HeatmapAnnotation object
  HeatmapAnnotation(df = anno,
                    col = anno_col,
                    show_annotation_name = TRUE)
}


## plot the distribution of certain statistic valuse. develop from statistics_plot
#' Plot the distribution of statistic valuse. develop from statistics_plot
#'
#' @param object A SummarizedExperiment or DEGdata obejct from \code{\link{test_diff}} / \code{\link{test_diff_deg}}
#' or \code{add_rejections}
#' @param statistic Character, in c("diff","t","p.val","p.adj")
#' @param contrasts Characters or NULL, plot contrasts. If NULL, plot all tested contrasts in objest.
#'
#' @return
#' The distribution histogram of statistics
#' @export
#'
#' @examples
#' # Load example
#' data(Silicosis_pg)
#' data <- Silicosis_pg
#' data_unique <- make_unique(data, "Gene.names", "Protein.IDs", delim = ";")
#'
#' # Differential test
#' ecols <- grep("LFQ.", colnames(data_unique))
#' se <- make_se_parse(data_unique, ecols,mode = "delim")
#' filt <- filter_se(se, thr = 0, fraction = 0.4, filter_formula = ~ Reverse != "+" & Potential.contaminant!="+")
#' norm <- normalize_vsn(filt)
#' imputed <- impute(norm, fun = "MinProb", q = 0.05)
#' diff <- test_diff(imputed, type = "control", control  = c("PBS"), fdr.type = "Storey's qvalue")
#' dep <- add_rejections(diff, alpha = 0.01,lfc = 2)
#'
#' # plot statistics distribution
#' plot_statistics(dep, statistic = "t.stastic")
#' plot_statistics(dep, statistic = "diff")
#'
plot_statistics <- function(object, statistic = c("diff","t.stastic","p.val","p.adj"), contrasts = NULL){

  assertthat::assert_that(all(statistic %in% c("diff","t.stastic","p.val","p.adj")),
                          class(object) %in% c("SummarizedExperiment", "DEGdata"))
  statistic = unique(statistic)
  col_data = colData(object)
  row_data = rowData(object)
  if (any(!c("label", "condition", "replicate") %in%
          colnames(col_data))) {
    stop(paste0("'label', 'condition' and/or 'replicate' columns are not present in '",
                deparse(substitute(object)), "'"), call. = FALSE)
  }
  if (length(grep("_diff", colnames(row_data))) < 1) {
    stop(paste0("'[contrast]_diff' columns are not present in '",
                deparse(substitute(object)), "'.\nRun test_diff() or test_diff_deg() to obtain the required columns."),
         call. = FALSE)
  }

  if(class(object) == "DEGdata"){
    if(any(statistic == "t.stastic")){
      warning("DEGdata draw a differetail expression test on DESeq2, t statistic do not exist in object.")
    }
    statistic = statistic[-which(statistic=="t.stastic")]
  }

  ## check contrasts
  exist_contrasts = get_contrast(object)
  if(is.null(contrasts)){
    contrasts = exist_contrasts
    # message("No input contrast and by_contrast is FALSE, use the foldchange in the first contrast: ",contrasts)
  }else if(any(!contrasts %in% exist_contrasts)){
    if(all(!contrasts %in% exist_contrasts)){
      stop("Input contrasts: ", paste0(contrasts, collapse = ", ")," don't exists. Contrasts should be: ",
           paste0(exist_contrasts, collapse = ", "))
    }else{
      warning("Input contrasts: ", paste0(contrasts[which(!contrasts %in% exist_contrasts)], collapse = ", ")," don't exists. Only perform on: ",
              paste0(intersect(contrasts, exist_contrasts), collapse = ", "))
      contrasts = intersect(contrasts, exist_contrasts)
    }
  }

  statistics_df = data.frame()
  # statistic = statistic
  # exist_statistics =
  # contrast = contrast
  statistics = paste0("_",statistic)
  for(i in contrasts){
    temp = row_data[,paste0(i, statistics),drop =F] %>% as.data.frame %>%
      mutate(contrast = i)
    colnames(temp)[-ncol(temp)] = colnames(temp)[-ncol(temp)] %>% sapply(.,function(x){strsplit(x,"_")[[1]] %>% tail(1)})
    statistics_df = rbind(statistics_df,temp)
  }

  p = list()
  for(i in statistic){
    p[[i]] = ggplot(statistics_df) +
      geom_histogram(aes_string(x = i), bins = 50)+
      facet_wrap(.~contrast) +
      labs( title = paste(i,"distribution"))+
      theme_DEP1()
  }

  if(length(p) == 1){
    return(p[[1]])
  }else if(length(p) > 1){
    p1 = p[[1]]
    for(i in 2:length(p)){
      p1 = p1/p[[i]]
    }
    return(p1)
  }else return(NULL)
}


statistics_plot <- function(dep,x = c("-log10.padj", "-log10.pval","padj","pval","t") ,
                            y = c("-log10.padj", "-log10.pval","padj","pval","t"), contrast,
                            plottype = c("x-y","histogram") ){

  row_data <- rowData(dep, use.names = FALSE) %>% as.data.frame()
  p.adjcol <- grep(paste("^",contrast, "_p.adj", sep = ""),colnames(row_data))
  p.valcol <- grep(paste("^",contrast, "_p.val", sep = ""),colnames(row_data))
  tcol <- grep(paste("^",contrast, "_t.stastic", sep = ""),colnames(row_data))
  sigcol <- grep(paste(contrast, "_significant", sep = ""),colnames(row_data))
  df = data.frame(log10.padj = -log10(row_data[,p.adjcol]), log10.pval = -log10(row_data[,p.valcol]),
                  padj = row_data[,p.adjcol], pval = row_data[,p.valcol],
                  t = row_data[,tcol], significant = row_data[,sigcol], assay(dep))

  df2 = data.frame(x = df[,gsub("-", "", x)], y = df[,gsub("-", "", y)])
  p <- ggplot(df2)
  if(plottype == "x-y"){
    p <- p + geom_point(aes(x,y)) +theme_bw() +
      xlab(label = x)+
      ylab(label = y)
  }

  if(plottype == "histogram"){
    p <- p + geom_histogram(mapping = aes(x),bins = 50) +
      theme_bw() +
      xlab(label = x)
  }
  p
  return(p)

}


#' Volcano plot
#'
#' \code{plot_volcano} generates a volcano plot for a specified contrast.
#'
#' @param object SummarizedExperiment,
#' Data object for which differentially enriched proteins are annotated
#' (output from \code{\link{test_diff}()} or \code{\link{test_diff_deg}()} and \code{\link{add_rejections}()}).
#' @param contrast Character(1),
#' Specifies the contrast to plot.
#' @param label_size numeric(1),
#' Sets the size of name labels.
#' @param add_names Logical(1),
#' Whether or not to plot names.
#' @param adjusted Logical(1),
#' Whether or not to use adjusted p values.
#' @param plot Logical(1),
#' If \code{TRUE} (default) the volcano plot is produced.
#' Otherwise (if \code{FALSE}), the data which the
#' volcano plot is based on are returned.
#' @param x_symmetry Logical(1),
#' if The X-axis is symmetric based on the 0.
#' @param breaks Numeric vector, the position of x-axis label.
#' eg: seq(-10,10,by = 5), c(-2, -1, 0, 1, 2).
#' @param label_number Integer(1).
#' The number of plot name. Defalut NULL, label all significant names.
#' @param chooseTolabel Character, Specify names to plot,
#' no matter whether they are significant.
#' @param onlyhighlight_Choose Logical(1), only highlight chosen labels.
#' Only work when parameter 'chooseTolabel' specify some feature names.
#' Features(proteins) in chooseTolabel will be highlighted by color, while
#' others will be colored in 'stable_color'.
#' @param label_trend Character(1).
#' one of c("all","up","down"), plot all significant, up-regulated ones or down-regulated ones.
#' @param add_threshold_line Character(1),
#' one of c("none", "intersect", "curve"). "intersect" draw a intersect cutoff line according
#' fcCutoff and pCutoff. "curve" draw a curve cutoff line according curvature and x0_fold.
#' @param fcCutoff numeric(1),
#' Cutoff of L2FC.
#' @param pCutoff numeric(1),
#' Cutoff of p or p adjusted value
#' @param curvature numeric(1),
#' Curvature of curve cutoff line
#' @param x0_fold numeric(1),
#' The fold of x0 to Sigma.
#' @param up_color The color of upregulated points
#' @param down_color The color of downregulated points
#' @param stable_color The color of stable points
#' @param highlight_PGs_with_few_peptides logic(1), mark the proteins with few peptides by \code{peptide_1_color}
#' and \code{peptide_2_color}. Require the peptide number variable(specified by \code{peptide_column}) in object.
#' @param peptide_1_color The color of points with 1 peptide.
#' @param peptide_2_color The color of points with 2 peptide.
#' @param peptide_column Character(1), the varable column name storing the peptide number.
#' @param max.overlaps A parameter from \code{\link[ggrepel]{geom_label_repel}}, which desides labels overlap when add_names is TRUE.
#' Exclude text labels when they overlap too many other things. Defaults to 30. Higher value allows more crowded labels.
#' @param dot_size numeric(1), the size of points.
#' @return A volcano plot (generated by \code{\link[ggplot2]{ggplot}})
#' @examples
#' # Load example
#' data(Silicosis_pg)
#' data <- Silicosis_pg
#' data_unique <- make_unique(data, "Gene.names", "Protein.IDs", delim = ";")
#'
#' # Differential test
#' ecols <- grep("LFQ.", colnames(data_unique))
#' se <- make_se_parse(data_unique, ecols,mode = "delim")
#' filt <- filter_se(se, thr = 0, fraction = 0.4, filter_formula = ~ Reverse != "+" & Potential.contaminant!="+")
#' norm <- normalize_vsn(filt)
#' imputed <- impute(norm, fun = "MinProb", q = 0.05)
#' diff <- test_diff(imputed, type = "control", control  = c("PBS"), fdr.type = "Storey's qvalue")
#' dep <- add_rejections(diff, alpha = 0.01,lfc = 2)
#'
#' plot_volcano(dep,contrast = "W4_vs_PBS")
#'
#' dep <- add_rejections(diff, thresholdmethod = "curve")
#' plot_volcano(dep,contrast = "W4_vs_PBS",add_threshold_line= "curve")
#' @export
plot_volcano <- function (object, contrast = get_contrast(object)[1],
                          add_names = TRUE,
                          dot_size = 1.5,
                          label_size = 3,
                          label_number = 40,
                          label_trend = c("all","up","down","none"),
                          up_color = "#B2182B",
                          down_color = "#2166AC",
                          stable_color = "#BEBEBE",
                          chooseTolabel = NULL,
                          onlyhighlight_Choose = FALSE,
                          adjusted = FALSE,
                          plot = TRUE, x_symmetry = TRUE,
                          add_threshold_line = c("none", "intersect", "curve"),
                          fcCutoff = 1, pCutoff = 0.05,
                          curvature = 0.6, x0_fold = 2,
                          breaks = NULL,
                          highlight_PGs_with_few_peptides = F,
                          peptide_1_color = "#A020F0", peptide_2_color = "#0000FF",
                          peptide_column = "Peptides",
                          max.overlaps = 30
){
  if (is.integer(label_size))
    label_size <- as.numeric(label_size)
  assertthat::assert_that(inherits(object, "SummarizedExperiment"),
                          is.character(contrast), length(contrast) == 1, is.numeric(label_size),
                          length(label_size) == 1, is.logical(add_names), length(add_names) ==
                            1, is.logical(adjusted), length(adjusted) == 1, is.logical(plot),
                          length(plot) == 1, is.logical(x_symmetry), length(x_symmetry) == 1,
                          is.character(label_trend), is.character(add_threshold_line),
                          is.numeric(fcCutoff),length(fcCutoff) == 1,
                          is.numeric(pCutoff),length(pCutoff) == 1,
                          is.numeric(curvature),length(curvature) == 1,
                          is.numeric(x0_fold),length(x0_fold) == 1
  )
  label_trend <- match.arg(label_trend)
  add_threshold_line <- match.arg(add_threshold_line)

  row_data <- rowData(object, use.names = TRUE)
  # if (any(!c("name", "ID") %in% colnames(row_data))) {
  #   stop(paste0("'name' and/or 'ID' columns are not present in '",
  #               deparse(substitute(object)), "'.\nRun make_unique() to obtain required columns."),
  #        call. = FALSE)
  # }
  if (length(grep("_p.adj|_diff", colnames(row_data))) < 1) {
    stop(paste0("'[contrast]_diff' and '[contrast]_p.adj' columns are not present in '",
                deparse(substitute(object)), "'.\nRun test_diff() to obtain the required columns."),
         call. = FALSE)
  }
  if (length(grep("_significant", colnames(row_data))) < 1 && add_threshold_line == "none") {
    stop(paste0("'[contrast]_significant' columns are not present in '",
                deparse(substitute(object)), "'.\nRun add_rejections() to obtain the required columns or set add_threshold_line."),
         call. = FALSE)
  }
  if (length(grep(paste("^", contrast, "_diff", sep = ""),
                  colnames(row_data))) == 0) {
    valid_cntrsts <- get_contrast(object)
    valid_cntrsts_msg <- paste0("Valid contrasts are: '",
                                paste0(valid_cntrsts, collapse = "', '"), "'")
    stop("Not a valid contrast, please run `plot_volcano()` with a valid contrast as argument\n",
         valid_cntrsts_msg, call. = FALSE)
  }


  #when name have Disadvantages eg: when contrast = c("H2A_vs_Biotin","uH2A_vs_Biotin"), when set contrast = c("H2A_vs_Biotin"), can find two cols,and then error, the same problem with p_values and signif
  #diff <- grep(paste(contrast, "_diff", sep = ""),
  #    colnames(row_data))
  diff <- match(paste(contrast, "_diff", sep = ""),
                colnames(row_data))
  if (adjusted) {
    p_values <- match(paste(contrast, "_p.adj", sep = ""),
                      colnames(row_data))
  }
  else {
    p_values <- match(paste(contrast, "_p.val", sep = ""),
                      colnames(row_data))
  }

  if (add_threshold_line == "none"){
    signif <- row_data[, match(paste(contrast, "_significant", sep = ""),
                               colnames(row_data))]
  }else if(add_threshold_line == "curve"){
    x = row_data[, diff]
    y = -log10(row_data[, p_values])
    x1 = fun.outlier(x) %>% na.omit()
    fit <- fitnormal(x1)
    Sigma = sqrt(fit$theta[2])
    x0 = x0_fold*Sigma

    curve_fun1 = function(x){
      y=abs(curvature/(x-x0))
    }
    curve_fun2 = function(x){
      y= abs(curvature/(x+x0))
    }
    signif <- ifelse(x>x0, y > curve_fun1(x),
                     ifelse(x < -x0,y > curve_fun2(x), FALSE))
  }else if(add_threshold_line == "intersect"){
    signif = (abs(row_data[, diff]) > fcCutoff & row_data[, p_values] < pCutoff)
  }


  df <- data.frame(x = row_data[, diff], y = -log10(row_data[, p_values]),
                   significant = signif, name = rownames(row_data)) %>%
    dplyr::filter(!is.na(significant)) %>% arrange(significant) %>%
    mutate(change = ifelse(significant, ifelse(x>0, 'up','down'), "stable"))

  # if
  if(onlyhighlight_Choose & length(chooseTolabel) > 0){
    df$change[which(! df$name %in% chooseTolabel)] = "stable"
    df$change[which(df$name %in% chooseTolabel)] =
      ifelse(df$x[which(df$name %in% chooseTolabel)] > 0,'up','down')
  }

  name1 <- gsub("_vs_.*", "", contrast)
  name2 <- gsub(".*_vs_", "", contrast)
  xlimit <- ceiling(max(abs(df$x)))

  # label_trend =
  #   down_color = up_color = "black"
  # stable_color = "grey"
  # stable_color = "#BEBEBE"
  if(label_trend == "none"){
    down_color <- up_color <- stable_color <- "black"
  }else if(label_trend == "up"){
    down_color = stable_color
  }else if(label_trend == "down"){
    up_color = stable_color
  }

  if(highlight_PGs_with_few_peptides ){
    if(peptide_column %in% colnames(row_data)){
      df$Peptides = row_data[match(df$name,rownames(row_data)),peptide_column] %>% as.numeric()
      df[!df$change == "stable",]$change[which(df[!df$change == "stable",]$Peptides == 1)] = "peptide_1"
      df[!df$change == "stable",]$change[which(df[!df$change == "stable",]$Peptides == 2)] = "peptide_2"
    }else{
      warning("peptide_column", peptide_column, "not exist in rowData of object. Please check")
    }
  }
  # df$change
  p <- ggplot(df, aes(x, y)) + geom_vline(xintercept = 0) +
    geom_point(aes(col = change),size = dot_size) +
    geom_text(data = data.frame(),
              aes(x = c(Inf, -Inf), y = c(-Inf, -Inf),
                  hjust = c(1,0), vjust = c(-1, -1),
                  label = c(name1, name2), size = 5, fontface = "bold")) +
    #scale_x_continuous(limits = c(-xlimit, xlimit), breaks = seq(-15, 15, by = 5)) +
    labs(title = contrast,
         x = expression(log[2] ~ "Fold change")) + theme_DEP1() +
    theme(legend.position = "none") + scale_color_manual(values = c(`up` = up_color,
                                                                    `down` = down_color,
                                                                    `stable` = stable_color,
                                                                    `peptide_1` = peptide_1_color,
                                                                    `peptide_2` = peptide_2_color))
  if (add_names) {
    df2 <- df %>% arrange(desc(significant),desc(y))
    # df2$label = df2$name
    df2$label = ""

    ## label label_number hits with highest p.
    if(!is.null(label_number) && is.null(chooseTolabel)) {
      df2$label[1:label_number] = df2$name[1:label_number]
    }

    ## label chooseTolabel no matter whether they are significant
    if(!is.null(chooseTolabel)){
      df2$label[which(df2$name %in% chooseTolabel)] = df2$name[which(df2$name %in% chooseTolabel)]
    }

    p <- p + ggrepel::geom_text_repel(data = dplyr::filter(df2, label != ""),
                                      aes(label = label, col = change), size = label_size,
                                      box.padding = unit(0.1, "lines"), point.padding = unit(0.2, "lines"),
                                      segment.size = 0.5,
                                      max.overlaps = max.overlaps)

  }
  if (adjusted) {
    p <- p + labs(y = expression(-log[10] ~ "Adjusted p-value"))
  }
  else {
    p <- p + labs(y = expression(-log[10] ~ "P-value"))

  }
  if(x_symmetry){
    p <- p + scale_x_continuous(limits = c(-xlimit, xlimit))
  }
  if(x_symmetry & !is.null(breaks)){
    p <- p + scale_x_continuous(limits = c(-xlimit, xlimit), breaks = breaks)
  }

  ## plot the cutoff line and cutoff parameters.
  linetype = 4
  if(add_threshold_line == "curve"){
    message("add curve threshold line." ," Sigma = ", Sigma, ", x0 = ", x0,", curvature = ", curvature )
    y_max <- layer_scales(p)$y$range$range[2] # store the y scale
    p <- p + stat_function(fun=curve_fun1,
                           xlim=c(x0,xlimit),
                           colour=up_color,linetype=linetype,lwd=0.6)
    p <- p + stat_function(fun=curve_fun2,
                           xlim=c(-xlimit,-x0),
                           colour=down_color,linetype=linetype,lwd=0.6)
    p <- p + ylim(0,y_max) # recover the y-axis scale
    p <- p + labs(subtitle =
                    bquote(
                      italic(log[10](p) ) > frac(italic(c), abs(italic(lfc)) - italic(x[0])) ~ "&" ~ italic(abs(lfc) > x0)
                      ~";"~italic(x[0]) == .(x0_fold) %*% .(round(Sigma,4))~italic(c) == .(round(curvature,3))
                    )
    ) + theme(plot.subtitle = element_text(hjust = 0.5))

  }else if(add_threshold_line == "intersect"){
    message("add intersect threshold line." ," fcCutoff = ", fcCutoff, "pCutoff = ", pCutoff)
    p <- p + geom_vline(xintercept=c(-fcCutoff,fcCutoff),linetype=linetype, col=c(down_color, up_color),lwd=0.6)
    p <- p + geom_hline(yintercept=-log10(pCutoff),linetype=linetype, col=stable_color,lwd=0.6)
    p <- p + labs(subtitle = bquote(
      fcCutoff==.(fcCutoff)~";"~pCutoff==.(pCutoff)
    ))+
      # paste0("fcCutoff=", fcCutoff, "; pCutoff=", pCutoff)) +
      theme(plot.subtitle = element_text(hjust = 0.5))
  }

  if (plot) {
    return(p)
  }
  else {
    df <- df %>% dplyr::select(name, x, y, significant, change) %>% arrange(desc(x))
    colnames(df)[c(1, 2, 3)] <- c("protein", "log2_fold_change",
                                  "p_value_-log10")
    if (adjusted) {
      colnames(df)[3] <- "adjusted_p_value_-log10"
    }
    return(df)
  }
}





#' Plot heatmap of specified genes/proteins across multiple omics results
#'
#' @param omics_list A list composed of SummarizedExperiment or DEGdata objects.
#' @param choose_name Character, the vector of specified identifiers.
#' @param color Character(1), sets the color panel (from \pkg{RColorBrewer}).
#' It can be a palette name in "RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral".
#' @param col_limit  Integer(1),
#' Sets the outer limits of the color scale.
#' @param to_upper Logical, whether transform all identifiers to upper
#' @param width Numeric, the width of heatmap
#' @param height Numeric, the height of heatmap
#' @param row_font_size Numeric, the size of row label.
#' @param ... other parameters to \code{\link[ComplexHeatmap]{Heatmap}}
#' @return
#' A HeatmapList from \code{ComplexHeatmap} package
#'
#' @export
#'
plot_multi_heatmap <- function(omics_list, choose_name, to_upper = FALSE,
                               color = c("RdBu", "RdYlBu", "RdYlGn", "BrBG", "PiYG", "PRGn", "PuOr", "RdGy", "Spectral"),
                               col_limit = 6,
                               width = 7, height = 10, row_font_size = 5, col_font_size = 11,
                               ...){
  assertthat::assert_that(class(omics_list) == "list", length(omics_list) > 0,
                          is.character(choose_name),
                          is.logical(to_upper), length(to_upper) == 1
                          )
  color <- match.arg(color)
  ht_list <- lapply(omics_list,function(x){
    if(class(x) == "DEGdata"){
      ht_mat = x@ntf
      rownames(ht_mat) = rownames(x)
    }else if(class(x) == "SummarizedExperiment"){
      ht_mat = assay(x)
    }
    if(to_upper)
      rownames(ht_mat) = rownames(ht_mat) %>% toupper() %>% make.names()
    return(ht_mat)
  })

  ht_list2 <- lapply(ht_list, function(x){
    x = x - rowMeans(x)
    x = as.data.frame(x)
    x = x[choose_name,]
    rownames(x) = choose_name
    x = as.matrix(x)
    return(x)
  })

  names(ht_list2) = names(omics_list)
  col_width = width/(sapply(ht_list2,ncol) %>% unlist %>% sum)

  heatmap_list <- lapply(1:length(ht_list2), function(x){
    mat = ht_list2[[x]]
    if(all(is.na(mat))){
      return(NULL)
    }else{
      ht = ComplexHeatmap::Heatmap(mat,
                                   heatmap_width = unit(col_width * ncol(mat) * 5, "cm"),
                                   # heatmap_height = unit(heatmap_height * 5, "cm"),
                                   col = circlize::colorRamp2(
                                     seq(-col_limit, col_limit, (col_limit/5)),
                                     rev(RColorBrewer::brewer.pal(11, color))),
                                   na_col = "grey80",
                                   cluster_rows = F,
                                   cluster_columns = F,
                                   column_title = names(ht_list2)[x],
                                   row_names_gp = gpar(fontsize = row_font_size),
                                   column_names_gp = gpar(fontsize = col_font_size),
                                   name = names(ht_list2)[x],
                                   heatmap_legend_param = list(title = names(ht_list2)[x] ),
                                   ...)
      return(ht)
    }
  })

  # heatmap_list[[1]] + heatmap_list[[2]] + heatmap_list[[3]] + heatmap_list[[4]]

  for (i in 1:length(heatmap_list)) {
    if(i == 1){
      heatmap_list2 = heatmap_list[[i]]
    }else{
      heatmap_list2 = heatmap_list2 + heatmap_list[[i]]
    }
  }

  heatmap_list2
  # draw(heatmap_list2,height = unit(height * 5, "cm"))
}


#' Plot venn plot of specified genes/proteins, or upon signicant results from
#'
#' Plot a venn plot for significant candidate in muitiple omics results, based on identifiers('name')
#'
#' @param x A list or a result SummarizedExperiment or DEGdata object.
#' If x is a list, it should be a list composed of vector or SummarizedExperiment or DEGdata,
#' plot_multi_venn compares vectors and significant result of SE objects in each list elements.
#' If x is a SummarizedExperiment or DEGdata object, it should be a result from add_rejections and contain more than than contrasts,
#' plot_multi_venn compares significant results in each contrasts.
#'
#' @param to_upper Logical, whether transform all identifiers to upper
#' @param background NULL or character vector of names of background protein/gene pool.
#' If background is provided, plot_multi_venn only consider candidates in background (the significant members in omics_list that also exist in background).
#' Else, if background is NULL, venn is ploted upon all candidates in omics_list.
#' @param plot  Logical(1), return a venn plot or a table.
#' @param contrasts NULL or vectors of characters. It only works when x is a SummarizedExperiment or DEGdata object from add_rejections.
#' To Specify compared contrasts in venn.
#' @inheritParams ggVennDiagram::ggVennDiagram
#' @param ... Additional arguments for imputation functions as depicted in \code{\link[ggVennDiagram]{ggVennDiagram}}.
#' @return
#' A venn plot by ggVennDiagram or a tibble (plot = F).
#' @note
#' Since DEP2 0.5.27, plot_multi_venn depends on ggVennDiagram (higher than 1.5).
#'
#' @export
#' @importFrom ggVennDiagram ggVennDiagram Venn process_data
#' @examples
#' data(Silicosis_pg)
#' data <- Silicosis_pg
#' data_unique <- make_unique(data, "Gene.names", "Protein.IDs", delim = ";")
#'
#' # Make SummarizedExperiment
#' ecols <- grep("LFQ.", colnames(data_unique))
#' se <- make_se_parse(data_unique, ecols, mode = "delim", sep = "_")
#'
#' filt <- filter_se(se, thr = 0, fraction = 0.3, filter_formula = ~ Reverse != "+" & Potential.contaminant!="+")
#' norm <- normalize_vsn(filt)
#' imputed <- impute(norm, fun = "MinDet", q = 0.01)
#'
#' # Test for differentially expressed proteins
#' diff <- test_diff(imputed, "control", "PBS", fdr.type = "BH")
#' dep <- add_rejections(diff, alpha = 0.01, lfc = 1)
#'
#' get_contrast(dep)
#'
#' # Plot venn of result
#' plot_multi_venn(
#'   dep,
#'   label_size = 2.5
#' )
#'
#' # Plot venn for some contrasts
#' get_contrast(dep)
#' plot_multi_venn(dep,contrasts = c("W10_vs_PBS", "W6_vs_PBS"), set_size = 3)
#'
#' # Compare with intersted geneset
#' IL1_relative_genes = c("Irg1", "Il1rn", "Saa3", "Zbp1", "Ccl6",
#'                        "Serpine1", "Ccl21a", "Pycard", "Irak2", "Vrk2",
#'                        "Fn1", "Il1r1", "Irf1", "Ccl9", "Mapk11", "Tank",
#'                        "Mapk13")
#' plot_multi_venn(
#'   list(W10 = get_signicant(dep,contrasts = "W10_vs_PBS"),
#'        W6 = get_signicant(dep,contrasts = "W6_vs_PBS"),
#'        IL1_genes = IL1_relative_genes
#'   )
#' )
#'
#' # Use background geneset
#' plot_multi_venn(
#'   list(W10 = get_signicant(dep,contrasts = "W10_vs_PBS"),
#'        W6 = get_signicant(dep,contrasts = "W6_vs_PBS"),
#'        IL1_genes = IL1_relative_genes
#'   ),
#'   background = IL1_relative_genes
#' )
#'
#' # Return a tibble
#' plot_multi_venn(
#'   dep,contrasts = c("W10_vs_PBS", "W6_vs_PBS"), plot = F,
#'   to_upper = T # transform all identifiers to upper
#' )
#'
plot_multi_venn <- function(x, to_upper = F, plot = T, contrasts = NULL,
                            background = NULL,
                            set_size = 5, label_size = 3, label_geom = "text",...){
  assertthat::assert_that(class(x) == "list" | inherits(x,"SummarizedExperiment"),
                          is.logical(to_upper), length(to_upper) == 1,
                          is.null(contrasts) | is.character(contrasts),
                          is.logical(plot), length(plot) == 1
  )

  if(inherits(x,"SummarizedExperiment")){
    # check SE contrasts.
    exist_contrasts = try({get_contrast(x)},silent = T)
    if(class(exist_contrasts)== "try-error"|length(exist_contrasts) < 1)
      stop("Check your input x! If x is a SummarizedExperiment or a DEGdata object, it should be an output from add_rejections!")
    if(length(exist_contrasts) == 1)
      stop(paste0("Check your input! Input x only has one contrast: ", contrasts,"."))

    # Used contrasts that are in exist_
    if(is.null(contrasts)){
      message(exist_contrasts)
      contrasts = exist_contrasts

    }else if(is.character(contrasts)){
      contrasts = intersect(contrasts, exist_contrasts)

      if(length(contrasts) < 1)
        stop(paste0("Check your input contrast! Only one contrast of no contrast is in ", deparse(substitute(diff))))

      message("Use contrasted: ",paste0(contrasts, collapse = ", "))
    }

    gene_list = lapply(contrasts,get_signicant, object = x, return_type = "names")
    names(gene_list) = contrasts
  }else if(is.list(x)){
    if(length(x) <= 1)
      stop("Check your input ",deparse(substitute(x)),". The x should have more than one elements.")

    # check list elements
    all_LISTorSE = x %>% lapply(function(y){
      is.character(y) | inherits(y,"SummarizedExperiment")
    }) %>% unlist
    if(!all(all_LISTorSE))
      stop("Check your input ",deparse(substitute(x)),". If x is a list, its elements should be a list composed of charecter vectors, SummarizedExperiment, or DEGdata objects")

    # extract significant results from SE elements
    gene_list = lapply(x, function(y){
      if(inherits(y,"SummarizedExperiment"))
        y = get_signicant(y, return_type = "names")
      return(y)
    })
  }

  if(to_upper){
    gene_list = lapply(gene_list, function(y){toupper(y)})
  }

  if(!is.null(background)){
    gene_list = lapply(gene_list, function(y){intersect(y,background)})
  }

  if(plot){
    ggp <- ggVennDiagram::ggVennDiagram(gene_list,
                                        set_size = set_size, label_size = label_size, label_geom = "text",
                                        ...)
    ggp = ggp + ggplot2::scale_x_continuous(expand = expansion(mult = .2))
    return(ggp)
  }else{
    venn <- ggVennDiagram::Venn(gene_list)
    data <- ggVennDiagram::process_data(venn)
    return(data$regionData)
  }
  # plotlyp <- ggp %>%
  #   ggplotly(tooltip = "text")

}



#' Fit a Gaussian distribution for L2FC of each contrast
#'
#' \code{plot_diff_hist} fit a Gaussian distribution for L2FC in each contrast.
#' The L2FC value is firstly filter out outlier, and then fit to normal curve based on likelihood.
#'
#' @param object a SummarizedExperiment or DEGdata output from \code{test_diff} or \code{test_diff_deg}
#' @param contrasts Character, the contrast in object
#' @param plot Logical(1), if TRUE return histogram(s) of L2FC in each contras with fitted normal curve.
#' if FALSE return the fitting result, mu and sigma of Gaussian distribution
#' @export
plot_diff_hist <- function(object, contrasts = NULL, plot = T){
  assertthat::assert_that(class(object) %in% c("SummarizedExperiment",
                                               "DEGdata"))
  # statistic = unique(statistic)
  col_data = colData(object)
  row_data = rowData(object)
  if (length(grep("_diff", colnames(row_data))) < 1) {
    stop(paste0("'[contrast]_diff' columns are not present in '",
                deparse(substitute(object)), "'.\nRun test_diff() or test_diff_deg() to obtain the required columns."),
         call. = FALSE)
  }

  exist_contrasts = get_contrast(object)
  if (is.null(contrasts)) {
    contrasts = exist_contrasts
  }else if (any(!contrasts %in% exist_contrasts)) {
    if (all(!contrasts %in% exist_contrasts)) {
      stop("Input contrasts: ", paste0(contrasts, collapse = ", "),
           " don't exists. Contrasts should be: ", paste0(exist_contrasts,
                                                          collapse = ", "))
    }
    else {
      warning("Input contrasts: ", paste0(contrasts[which(!contrasts %in%
                                                            exist_contrasts)], collapse = ", "), " don't exists. Only perform on: ",
              paste0(intersect(contrasts, exist_contrasts),
                     collapse = ", "))
      contrasts = intersect(contrasts, exist_contrasts)
    }
  }

  ## get l2fc data
  statistics_df = row_data[,paste0(contrasts,"_diff"),drop = F] %>% data.frame()
  ## fit sigma and mu of normal distribution for each contrast
  fit_res <- apply(statistics_df, 2, function(x){
    x = fun.outlier(as.vector(x),time.iqr = 1.5) %>% na.omit()
    fit = fitnormal(x)
    return(fit$theta)
  })
  fit_res = fit_res %>% t %>% data.frame %>% rownames_to_column(var = "contrast")

  if(plot){ ## plot hist with gaussian curve
    statistics_df <- statistics_df %>% rownames_to_column(var = "name") %>%
      gather(., "contrasts",value = "diff", -name) %>%
      mutate(contrasts = gsub("_diff$","",contrasts))

    ps <- lapply(contrasts, function(thecontrasts){
      the_statistics_df = statistics_df %>% filter(contrasts == thecontrasts)
      the_fit_res = fit_res  %>% filter(contrasts == thecontrasts)
      ggplot(the_statistics_df) +
        geom_histogram(aes(x = diff, y = after_stat(density)), bins = 60)+
        geom_function(fun = dnorm, args = list(mean = the_fit_res$mu, sd = sqrt(the_fit_res$sigma2)), colour = "blue3") +
        labs(title = paste("L2FC", thecontrasts),
             subtitle = bquote(italic(N)(mu == .(round(the_fit_res$mu,5)), sigma^2 == .(round(the_fit_res$sigma2,5)) ))) +
        theme_DEP2() + theme(plot.subtitle = element_text(hjust=1))
      # +xlim(quantile(the_statistics_df$diff,0.25) - 4*IQR(the_statistics_df$diff),
      #        quantile(the_statistics_df$diff,0.75) + 4*IQR(the_statistics_df$diff))
    })

    if(length(ps) > 1){
      return(multiplot(plotlist  = ps, cols=2))
    }else{ ## one contrast
      return(ps[[1]]) }
  }else{
    return(fit_res)
  }
}

# Multiple plot function from cookbook-r
#
# ggplot objects can be passed in ..., or to plotlist (as a list of ggplot objects)
# - cols:   Number of columns in layout
# - layout: A matrix specifying the layout. If present, 'cols' is ignored.
#
# If the layout is something like matrix(c(1,2,3,3), nrow=2, byrow=TRUE),
# then plot 1 will go in the upper left, 2 will go in the upper right, and
# 3 will go all the way across the bottom.
#
multiplot <- function(..., plotlist=NULL, file, cols=1, layout=NULL) {
  library(grid)

  # Make a list from the ... arguments and plotlist
  plots <- c(list(...), plotlist)

  numPlots = length(plots)

  # If layout is NULL, then use 'cols' to determine layout
  if (is.null(layout)) {
    # Make the panel
    # ncol: Number of columns of plots
    # nrow: Number of rows needed, calculated from # of cols
    layout <- matrix(seq(1, cols * ceiling(numPlots/cols)),
                     ncol = cols, nrow = ceiling(numPlots/cols))
  }

  if (numPlots==1) {
    print(plots[[1]])

  } else {
    # Set up the page
    grid.newpage()
    pushViewport(viewport(layout = grid.layout(nrow(layout), ncol(layout))))

    # Make each plot, in the correct location
    for (i in 1:numPlots) {
      # Get the i,j matrix positions of the regions that contain this subplot
      matchidx <- as.data.frame(which(layout == i, arr.ind = TRUE))

      print(plots[[i]], vp = viewport(layout.pos.row = matchidx$row,
                                      layout.pos.col = matchidx$col))
    }
  }
}
